import numpy as np
import os
import joblib
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from config import BASE_DIR, OUTPUT_CONFIG
import matplotlib as mpl
from datetime import datetime

# 设置Agg后端，确保在无GUI环境下也能正常生成图像
mpl.use('Agg')

# 禁用中文字体设置，使用英文标签避免乱码
# mpl.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
# mpl.rcParams['axes.unicode_minus'] = False

def load_model():
    """加载训练好的模型和多项式特征转换器，如果失败则创建新模型"""
    model_path = os.path.join(BASE_DIR, OUTPUT_CONFIG['model_file'])
    poly_path = os.path.join(BASE_DIR, OUTPUT_CONFIG['poly_file'])
    
    try:
        model = joblib.load(model_path)
        poly = joblib.load(poly_path)
        print(f"已加载模型: {model_path}")
        
        # 简单测试模型
        test_speed = 0.4
        test_array = np.array([[test_speed]])
        test_poly = poly.transform(test_array)
        test_pred = model.predict(test_poly)[0]
        
        if test_pred < 600 or test_pred > 1900:
            print(f"警告: 模型测试不通过，预测值异常: {test_pred}，将重新训练模型")
            raise Exception("模型预测异常")
            
        return model, poly
    except Exception as e:
        print(f"加载模型失败或模型异常: {e}")
        print("将创建新的预测模型")
        return None, None

def load_calibration_data():
    """加载标定数据，用于参考和验证"""
    csv_path = os.path.join(BASE_DIR, OUTPUT_CONFIG['combined_results_file'])
    
    if not os.path.exists(csv_path):
        print(f"找不到标定数据文件: {csv_path}")
        return None
        
    try:
        data = pd.read_csv(csv_path)
        print(f"已加载标定数据: {csv_path}")
        print(f"数据范围: 速度 {data['SteadyStateSpeed'].min():.4f}-{data['SteadyStateSpeed'].max():.4f} m/s, 指令 {data['Command'].min():.0f}-{data['Command'].max():.0f}")
        return data
    except Exception as e:
        print(f"加载标定数据失败: {e}")
        return None

def find_closest_in_data(speed, data):
    """在实际标定数据中查找最接近的速度对应的命令值"""
    if data is None:
        return None, None
        
    # 计算输入速度与所有标定速度的差值绝对值
    data_copy = data.copy()
    data_copy['Diff'] = abs(data_copy['SteadyStateSpeed'] - speed)
    
    # 找到差值最小的行
    closest = data_copy.loc[data_copy['Diff'].idxmin()]
    
    return closest['Command'], closest['SteadyStateSpeed']

def find_closest_command(command, data):
    """在实际标定数据中查找最接近的命令对应的速度值"""
    if data is None:
        return None, None
        
    # 计算输入命令与所有标定命令的差值绝对值
    data_copy = data.copy()
    data_copy['Diff'] = abs(data_copy['Command'] - command)
    
    # 找到差值最小的行
    closest = data_copy.loc[data_copy['Diff'].idxmin()]
    
    return closest['Command'], closest['SteadyStateSpeed']

def train_new_model(data):
    """训练速度到命令的新模型"""
    if data is None:
        print("错误: 无法训练模型，没有标定数据")
        return None, None
    
    # 提取特征和目标
    X = data[['SteadyStateSpeed']].values
    y = data['Command'].values
    
    # 创建多项式特征 (较低阶多项式可能更稳定)
    poly = PolynomialFeatures(degree=2)
    X_poly = poly.fit_transform(X)
    
    # 训练线性回归模型
    model = LinearRegression()
    model.fit(X_poly, y)
    
    # 保存模型
    model_path = os.path.join(BASE_DIR, "new_speed_model.joblib")
    poly_path = os.path.join(BASE_DIR, "new_poly_features.joblib")
    joblib.dump(model, model_path)
    joblib.dump(poly, poly_path)
    print(f"已保存新模型到: {model_path}")
    
    return model, poly

def predict_command(speed, model, poly, data=None, prefer_data=False):
    """使用模型预测给定速度对应的指令值"""
    if speed <= 0:
        return 0
    
    # 方法1: 直接从数据查找最近点（如果可用且优先）
    if data is not None and prefer_data:
        nearest_command, _ = find_closest_in_data(speed, data)
        # 如果找到了数据点，直接返回
        if nearest_command is not None:
            return nearest_command
    
    # 方法2: 使用线性插值（如果有数据但速度在范围内）
    if data is not None and 0.05 < speed < 1.0:
        try:
            # 对数据按速度排序
            sorted_data = data.sort_values('SteadyStateSpeed')
            
            # 如果速度小于最小值或大于最大值，使用最近点
            if speed <= sorted_data['SteadyStateSpeed'].min():
                return sorted_data.iloc[0]['Command']
            if speed >= sorted_data['SteadyStateSpeed'].max():
                return sorted_data.iloc[-1]['Command']
            
            # 否则进行线性插值
            for i in range(len(sorted_data) - 1):
                speed1 = sorted_data.iloc[i]['SteadyStateSpeed']
                speed2 = sorted_data.iloc[i+1]['SteadyStateSpeed']
                
                if speed1 <= speed <= speed2:
                    cmd1 = sorted_data.iloc[i]['Command']
                    cmd2 = sorted_data.iloc[i+1]['Command']
                    
                    # 线性插值公式
                    ratio = (speed - speed1) / (speed2 - speed1)
                    interpolated_cmd = cmd1 + ratio * (cmd2 - cmd1)
                    
                    # 四舍五入到整数
                    return round(interpolated_cmd)
        except Exception as e:
            print(f"线性插值失败: {e}，使用模型预测")
    
    # 方法3: 使用多项式模型（如果可用）
    if model is not None and poly is not None:
        # 将速度转换为二维数组
        speed_array = np.array([[speed]])
        
        # 使用多项式特征转换
        speed_poly = poly.transform(speed_array)
        
        # 预测指令值
        command = model.predict(speed_poly)[0]
        
        # 确保命令值在合理范围内 (500-2000)
        command = max(500, min(2000, command))
        
        return command
    
    # 方法4: 回退到简单线性映射
    if data is not None:
        min_cmd = data['Command'].min()
        max_cmd = data['Command'].max()
        min_speed = data['SteadyStateSpeed'].min()
        max_speed = data['SteadyStateSpeed'].max()
        
        # 线性映射公式
        if max_speed > min_speed:  # 防止除以零
            cmd_range = max_cmd - min_cmd
            speed_range = max_speed - min_speed
            ratio = cmd_range / speed_range
            command = min_cmd + (speed - min_speed) * ratio
            
            # 确保命令值在合理范围内
            command = max(500, min(2000, command))
            return command
    
    # 如果以上方法都不可用，返回合理的默认值
    return 1500

def predict_speed(command, data=None, inverse_model=None, inverse_poly=None):
    """根据指令值预测速度"""
    if command < 500 or command > 2000:
        print(f"警告: 指令值 {command} 超出范围 (500-2000)，已调整")
        command = max(500, min(2000, command))
    
    # 方法1: 直接从数据查找最近点
    if data is not None:
        _, closest_speed = find_closest_command(command, data)
        # 如果找到了数据点，直接返回
        if closest_speed is not None:
            return closest_speed
    
    # 方法2: 使用线性插值（如果有数据）
    if data is not None:
        try:
            # 对数据按命令值排序
            sorted_data = data.sort_values('Command')
            
            # 如果命令值小于最小值或大于最大值，使用最近点
            if command <= sorted_data['Command'].min():
                return sorted_data.iloc[0]['SteadyStateSpeed']
            if command >= sorted_data['Command'].max():
                return sorted_data.iloc[-1]['SteadyStateSpeed']
            
            # 否则进行线性插值
            for i in range(len(sorted_data) - 1):
                cmd1 = sorted_data.iloc[i]['Command']
                cmd2 = sorted_data.iloc[i+1]['Command']
                
                if cmd1 <= command <= cmd2:
                    speed1 = sorted_data.iloc[i]['SteadyStateSpeed']
                    speed2 = sorted_data.iloc[i+1]['SteadyStateSpeed']
                    
                    # 线性插值公式
                    ratio = (command - cmd1) / (cmd2 - cmd1)
                    interpolated_speed = speed1 + ratio * (speed2 - speed1)
                    
                    return interpolated_speed
        except Exception as e:
            print(f"线性插值失败: {e}")
    
    # 方法3: 使用反向模型（如果可用）
    if inverse_model is not None and inverse_poly is not None:
        try:
            # 将命令转换为二维数组
            cmd_array = np.array([[command]])
            
            # 使用多项式特征转换
            cmd_poly = inverse_poly.transform(cmd_array)
            
            # 预测速度值
            speed = inverse_model.predict(cmd_poly)[0]
            
            # 确保速度值在合理范围内
            speed = max(0.05, min(2.0, speed))
            
            return speed
        except Exception as e:
            print(f"反向模型预测失败: {e}")
    
    # 方法4: 回退到简单线性映射
    if data is not None:
        min_cmd = data['Command'].min()
        max_cmd = data['Command'].max()
        min_speed = data['SteadyStateSpeed'].min()
        max_speed = data['SteadyStateSpeed'].max()
        
        # 线性映射公式
        cmd_range = max_cmd - min_cmd
        if cmd_range > 0:  # 防止除以零
            speed_range = max_speed - min_speed
            ratio = speed_range / cmd_range
            speed = min_speed + (command - min_cmd) * ratio
            
            # 确保速度值在合理范围内
            speed = max(0.05, min(1.0, speed))
            return speed
    
    # 如果以上方法都不可用，返回合理的默认值
    return 0.3

def load_or_create_inverse_model(data):
    """加载或创建反向模型（从指令到速度）"""
    model_path = os.path.join(BASE_DIR, "new_inverse_model.joblib")
    poly_path = os.path.join(BASE_DIR, "new_inverse_poly.joblib")
    
    if os.path.exists(model_path) and os.path.exists(poly_path):
        try:
            model = joblib.load(model_path)
            poly = joblib.load(poly_path)
            print(f"已加载反向模型: {model_path}")
            return model, poly
        except Exception as e:
            print(f"加载反向模型失败: {e}")
    else:
        print(f"未找到反向模型，尝试从数据创建")
    
    # 如果没有预训练模型，从数据创建简单模型
    if data is not None:
        # 提取特征和目标
        X = data[['Command']].values
        y = data['SteadyStateSpeed'].values
        
        # 创建多项式特征
        poly = PolynomialFeatures(degree=2)
        X_poly = poly.fit_transform(X)
        
        # 训练线性回归模型
        model = LinearRegression()
        model.fit(X_poly, y)
        
        # 保存模型
        joblib.dump(model, model_path)
        joblib.dump(poly, poly_path)
        print(f"已创建并保存新的反向模型")
        
        return model, poly
    
    return None, None

def plot_prediction_range(model, poly, inverse_model, inverse_poly, data=None):
    """绘制预测范围图表"""
    # 设置matplotlib后端为Agg
    mpl.use('Agg')
    
    plt.figure(figsize=(12, 8))
    
    # 获取当前时间作为时间戳
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 设置更大的预测范围
    speed_min = 0.0
    speed_max = 2.0
    cmd_min = 0
    cmd_max = 2500
    
    # 添加实际数据点（如果有）
    if data is not None:
        unique_datasets = data['Dataset'].unique()
        colors = ['green', 'orange', 'purple', 'brown', 'pink']
        markers = ['o', 's', '^', 'x', 'd']
        
        for i, dataset in enumerate(unique_datasets):
            dataset_data = data[data['Dataset'] == dataset]
            color = colors[i % len(colors)]
            marker = markers[i % len(markers)]
            plt.scatter(dataset_data['SteadyStateSpeed'], dataset_data['Command'], 
                       label=f'Data: {dataset}', color=color, marker=marker, s=50)
    
    # 绘制速度到指令的预测曲线（使用插值方法）
    speed_range = np.linspace(speed_min, speed_max, 200)
    cmd_preds_model = []
    cmd_preds_interp = []
    
    for speed in speed_range:
        # 模型预测
        if model is not None and poly is not None:
            cmd_model = predict_command(speed, model, poly, data, prefer_data=False)
            cmd_preds_model.append(cmd_model)
        
        # 插值预测
        cmd_interp = predict_command(speed, None, None, data, prefer_data=True)
        cmd_preds_interp.append(cmd_interp)
    
    if model is not None and poly is not None:
        plt.plot(speed_range, cmd_preds_model, 'r-', linewidth=2, label='Model Prediction')
    
    plt.plot(speed_range, cmd_preds_interp, 'g--', linewidth=2, label='Interpolation Prediction')
    
    # 添加网格和标签
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.xlabel('Speed (m/s)', fontsize=12)
    plt.ylabel('Command Value', fontsize=12)
    plt.title(f'Speed-Command Prediction Relationship\n{current_time}', fontsize=14)
    
    # 设置坐标轴范围
    plt.xlim(speed_min, speed_max)
    plt.ylim(cmd_min, cmd_max)
    
    # 添加标注线
    plt.axhline(y=500, color='gray', linestyle='--', alpha=0.5)
    plt.axhline(y=2000, color='gray', linestyle='--', alpha=0.5)
    plt.axvline(x=0.1, color='gray', linestyle='--', alpha=0.5)
    plt.axvline(x=0.5, color='gray', linestyle='--', alpha=0.5)
    
    # 添加区域标注（使用英文）
    plt.text(0.12, 450, 'Calibration Range', fontsize=10, color='gray')
    plt.text(0.55, 450, 'Prediction Range (Uncertain Area)', fontsize=10, color='gray')
    
    plt.legend(loc='best')
    
    # 保存图表（使用时间戳和英文文件名）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    plot_path = os.path.join(BASE_DIR, f"speed_command_prediction_{timestamp}.png")
    plt.savefig(plot_path, dpi=300)
    print(f"Prediction range plot saved to: {plot_path}")
    
    plt.close()

def main():
    # 加载标定数据 (首先加载数据，因为模型训练需要数据)
    data = load_calibration_data()
    
    if data is None:
        print("错误: 无法加载标定数据，请先运行数据分析.py生成数据")
        return
    
    # 加载或创建模型
    model, poly = load_model()
    
    # 如果模型加载失败，尝试创建新模型
    if model is None or poly is None:
        model, poly = train_new_model(data)
    
    # 加载或创建反向模型
    inverse_model, inverse_poly = load_or_create_inverse_model(data)
    
    # 生成预测范围图表
    plot_prediction_range(model, poly, inverse_model, inverse_poly, data)
    
    print("\n速度-指令预测系统 (改进版)")
    print("输入范围: 0.0 - 2.0 m/s (推荐: 0.1 - 0.5 m/s)")
    print("特殊命令:")
    print("  'q' - 退出程序")
    print("  'i' - 切换到指令→速度模式")
    print("  's' - 切换到速度→指令模式")
    print("  'm' - 切换预测方法 (模型/插值)")
    print("  'r' - 输入一系列速度值")
    print("  't' - 生成预测表格")
    
    # 默认模式: 速度→指令
    mode = "speed_to_cmd"
    # 默认使用插值方法
    use_model = False
    
    while True:
        try:
            if mode == "speed_to_cmd":
                # 速度→指令模式
                value_input = input("\n请输入目标速度 (m/s): ")
            else:
                # 指令→速度模式
                value_input = input("\n请输入指令值 (500-2000): ")
            
            # 处理特殊命令
            if value_input.lower() == 'q':
                break
            elif value_input.lower() == 'i':
                mode = "cmd_to_speed"
                print("已切换到指令→速度模式")
                continue
            elif value_input.lower() == 's':
                mode = "speed_to_cmd"
                print("已切换到速度→指令模式")
                continue
            elif value_input.lower() == 'm':
                use_model = not use_model
                if use_model:
                    print("已切换到模型预测方法")
                else:
                    print("已切换到数据插值预测方法")
                continue
            elif value_input.lower() == 'r' and mode == "speed_to_cmd":
                generate_range_predictions(model, poly, inverse_model, inverse_poly, data, use_model)
                continue
            elif value_input.lower() == 't':
                generate_prediction_table(model, poly, inverse_model, inverse_poly, data, use_model)
                continue
            
            # 转换为浮点数
            value = float(value_input)
            
            if mode == "speed_to_cmd":
                # 速度→指令模式
                # 检查输入速度并警告
                if value < 0:
                    print("警告: 输入速度小于0，已调整为0")
                    value = 0
                elif value > 2.0:
                    print("警告: 输入速度超过2.0 m/s，预测结果可能不准确")
                elif value > 0.5:
                    print("警告: 输入速度超出标定范围 (0.1-0.5 m/s)，预测结果可能不准确")
                
                # 预测指令值
                if use_model and model is not None and poly is not None:
                    # 使用模型预测
                    pred_command = predict_command(value, model, poly, data, prefer_data=False)
                    method_name = "模型"
                else:
                    # 使用插值预测
                    pred_command = predict_command(value, None, None, data, prefer_data=True)
                    method_name = "插值"
                
                # 输出预测结果
                print(f"预测的指令值 ({method_name}方法): {pred_command:.0f}")
                
                # 验证预测结果
                expected_speed = predict_speed(pred_command, data, inverse_model, inverse_poly)
                print(f"预期实际速度: {expected_speed:.4f} m/s (误差: {abs(expected_speed-value):.4f} m/s)")
                
                # 显示最接近的实际标定点
                closest_cmd, closest_speed = find_closest_in_data(value, data)
                if closest_cmd is not None:
                    print(f"最接近的标定点: 速度={closest_speed:.4f} m/s, 指令={closest_cmd:.0f}")
            else:
                # 指令→速度模式
                # 预测速度值
                pred_speed = predict_speed(value, data, inverse_model, inverse_poly)
                print(f"预测的速度: {pred_speed:.4f} m/s")
                
                # 验证预测
                if use_model and model is not None and poly is not None:
                    # 使用模型验证
                    cmd_verify = predict_command(pred_speed, model, poly, data, prefer_data=False)
                    method_name = "模型"
                else:
                    # 使用插值验证
                    cmd_verify = predict_command(pred_speed, None, None, data, prefer_data=True)
                    method_name = "插值"
                    
                print(f"验证指令值 ({method_name}方法): {cmd_verify:.0f} (误差: {abs(cmd_verify-value):.0f})")
                
                # 显示最接近的标定点
                closest_cmd, closest_speed = find_closest_command(value, data)
                if closest_cmd is not None:
                    print(f"最接近的标定点: 指令={closest_cmd:.0f}, 速度={closest_speed:.4f} m/s")
            
        except ValueError:
            print("请输入有效的数字")
        except Exception as e:
            print(f"发生错误: {e}")
            import traceback
            traceback.print_exc()

def generate_range_predictions(model, poly, inverse_model, inverse_poly, data, use_model=False):
    """生成一系列速度值对应的指令预测"""
    try:
        min_val = float(input("请输入起始速度 (m/s): "))
        max_val = float(input("请输入结束速度 (m/s): "))
        step = float(input("请输入步长 (m/s): "))
        
        if min_val > max_val:
            min_val, max_val = max_val, min_val
        
        if step <= 0:
            step = 0.1
            
        print("\n速度 (m/s) | 指令值 | 验证速度 (m/s) | 误差 (m/s)")
        print("-" * 50)
        
        # 获取当前时间作为时间戳
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 准备保存结果
        results = []
        
        speed_vals = np.arange(min_val, max_val + step, step)
        for speed in speed_vals:
            # 根据选择的方法预测指令值
            if use_model and model is not None and poly is not None:
                cmd = predict_command(speed, model, poly, data, prefer_data=False)
                method = "模型"
                method_eng = "model"
            else:
                cmd = predict_command(speed, None, None, data, prefer_data=True)
                method = "插值"
                method_eng = "interpolation"
            
            # 验证
            verify_speed = predict_speed(cmd, data, inverse_model, inverse_poly)
            error = abs(verify_speed - speed)
            print(f"{speed:.2f} | {cmd:.0f} | {verify_speed:.4f} | {error:.4f}")
            
            # 添加到结果列表
            results.append([speed, cmd, verify_speed, error, method])
        
        # 保存结果到CSV文件
        save_option = input("\n是否保存结果到CSV文件? (y/n): ")
        if save_option.lower() == 'y':
            df = pd.DataFrame(results, columns=["目标速度(m/s)", "预测指令值", "验证速度(m/s)", "误差(m/s)", "预测方法"])
            output_path = os.path.join(BASE_DIR, f"speed_command_range_{min_val:.2f}to{max_val:.2f}_{method_eng}_{current_time}.csv")
            df.to_csv(output_path, index=False)
            print(f"预测范围结果已保存到: {output_path}")
        
    except ValueError:
        print("请输入有效的数字")
    except Exception as e:
        print(f"生成预测范围时出错: {e}")

def generate_prediction_table(model, poly, inverse_model, inverse_poly, data, use_model=False):
    """生成预测表格并保存为CSV文件"""
    try:
        # 获取当前时间作为时间戳
        current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 速度范围: 0.05到2.0, 步长0.05
        speeds = np.arange(0.05, 2.01, 0.05)
        results = []
        
        for speed in speeds:
            # 根据选择的方法预测指令值
            if use_model and model is not None and poly is not None:
                cmd = predict_command(speed, model, poly, data, prefer_data=False)
                method = "模型"
                method_eng = "model"
            else:
                cmd = predict_command(speed, None, None, data, prefer_data=True)
                method = "插值"
                method_eng = "interpolation"
            
            # 验证
            verify_speed = predict_speed(cmd, data, inverse_model, inverse_poly)
            error = abs(verify_speed - speed)
            results.append([speed, cmd, verify_speed, error, method])
        
        # 创建DataFrame
        df = pd.DataFrame(results, columns=["目标速度(m/s)", "预测指令值", "验证速度(m/s)", "误差(m/s)", "预测方法"])
        
        # 保存为CSV - 使用带时间戳的文件名
        output_path = os.path.join(BASE_DIR, f"speed_command_table_{method_eng}_{current_time}.csv")
        df.to_csv(output_path, index=False)
        print(f"预测表格已保存到: {output_path}")
        
        # 打印一部分结果
        print("\n预测表格摘要 (每0.2 m/s一个样本):")
        print(df.iloc[::4].to_string(index=False))
        
    except Exception as e:
        print(f"生成预测表格时出错: {e}")

if __name__ == "__main__":
    main() 